#include <thrust/copy.h>
#include <cusp/precond/spai.h>
#include <cusp/krylov/bicgstab.h>
#include <cusp/gallery/poisson.h>
#include <cusp/gallery/eye.h>
#include <cusp/csr_matrix.h>

#include <cusp/frobenius_norm.h>
#include <cusp/print.h>
#include <cusp/multiply.h>
#include <cusp/elementwise.h>
#include <cusp/csr_matrix.h>
#include <cusp/graph/neighborhood_extension.h>

#include <iostream>

#define NEIGH_LEV 4

// Number of levels of neighbors to be considered for considered

template <typename Monitor>
void report_status(Monitor& monitor)
{
    if (monitor.converged())
    {
        std::cout << "Solver converged to " << monitor.tolerance() << " tolerance";
        std::cout << " after " << monitor.iteration_count() << " iterations";
        std::cout << " (" << monitor.residual_norm() << " final residual)" << std::endl;
    }
    else
    {
        std::cout << "Solver reached iteration limit " << monitor.iteration_limit() << " before converging";
        std::cout << " to " << monitor.tolerance() << " tolerance ";
        std::cout << " (" << monitor.residual_norm() << " final residual)" << std::endl;
    }
}

int main(int argc, char** argv)
{
    typedef int                 IndexType;
    typedef double              ValueType;
    typedef cusp::device_memory MemorySpace;

    // create an empty sparse matrix structure
    cusp::csr_matrix<IndexType, ValueType, cusp::host_memory> A_csr;
    cusp::array2d<ValueType, MemorySpace> rhs(A_csr.num_rows, 1);
    cusp::csr_matrix<IndexType, ValueType, cusp::host_memory> rhs_csr;

    if (argc < 2)
    {
        cusp::gallery::poisson5pt(A_csr, 512, 512);
    }
    else
    {
        // read matrix from file
        cusp::io::read_matrix_market_file(A_csr, std::string(argv[1]));

        if (argc > 2)  // read right-hand side
        {
            cusp::io::read_matrix_market_file(rhs, std::string(argv[2]));
        }
    }

    std::cout << "size " << A_csr.num_rows << " x " << A_csr.num_cols << std::endl;

    cusp::csr_matrix<IndexType, ValueType, MemorySpace> d_A(A_csr);

//    cusp::print(SP_csr);

    // solve without preconditioning
    {
        std::cout << "\nSolving with no preconditioner" << std::endl;
    
        // allocate storage for solution (x) and right hand side (b)
        cusp::array1d<ValueType, MemorySpace> x(d_A.num_rows, 0);
        cusp::array1d<ValueType, MemorySpace> b(d_A.num_rows, 1);
        if (argc > 2)  // use the rhs from the file
           thrust::copy(rhs.values.begin(),rhs.values.end(),b.begin());
        // cusp::print(b);
        // set stopping criteria (iteration_limit = 1000, relative_tolerance = 1e-6)
        cusp::verbose_monitor<ValueType> monitor(b, 1000, 1e-6);
        
        // solve
        cusp::krylov::bicgstab(d_A, x, b, monitor);

        // report status
        report_status(monitor);
    }


    // solve SPAI preconditioner
    for (IndexType lev=0; lev< NEIGH_LEV; lev++)
    {
        cusp::csr_matrix<IndexType, ValueType, cusp::host_memory> SP_csr(cusp::graph::neighborhood_extension(A_csr,lev)); // Copy constructor

        

        cusp::csr_matrix<IndexType, ValueType, MemorySpace> d_SP(SP_csr);
        std::cout << "\nSolving with SPAI preconditioner with neighbor level " << lev << " non-zeros: " << SP_csr.num_entries << std::endl;
        
        // allocate storage for solution (x) and right hand side (b)
        cusp::array1d<ValueType, MemorySpace> x(d_A.num_rows, 0);
        cusp::array1d<ValueType, MemorySpace> b(d_A.num_rows, 1);
        thrust::copy(rhs_csr.values.begin(),rhs_csr.values.end(),b.begin());

        // set stopping criteria (iteration_limit = 1000, relative_tolerance = 1e-6)
        cusp::verbose_monitor<ValueType> monitor(b, 1000, 1e-6);

        // setup preconditioner
	cusp::csr_matrix<IndexType, ValueType, MemorySpace> I;
	cusp::gallery::eye(I, d_A.num_rows, d_A.num_cols);
        cusp::precond::spai<ValueType, MemorySpace> M(d_A, d_SP);
        std::cout << "Completed preconditioner \n";

        // solve
        cusp::krylov::bicgstab(d_A, x, b, monitor, M);
        
        // report status
        report_status(monitor);

        // report norms
        cusp::csr_matrix<int,double,cusp::device_memory> P;
        cusp::multiply(d_A, M.matrix, P);
        cusp::subtract(P, I, P);
        std::cout << "Fnorm(A)\t" << cusp::frobenius_norm(d_A) << "\n";
        std::cout << "Fnorm(AM-I)\t" << cusp::frobenius_norm(P) << std::endl;
    }

    return 0;
}

